package model

import (
	"bufio"
	"fmt"
	"gitee.com/swsk33/shard-download-demo/util"
	"github.com/fatih/color"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

// ParallelGetTask 多线程下载任务类，存放一个多线程下载任务的参数和状态量
type ParallelGetTask struct {
	// 文件的下载链接
	Url string
	// 文件的最终保存位置
	FilePath string
	// 下载并发数
	Concurrent int
	// 下载的分片临时文件保存文件夹
	TempFolder string
	// 下载文件的总大小
	TotalSize int64
	// 全部的下载分片任务参数列表
	ShardTaskList []*ShardTask
}

// NewParallelGetTask 构造函数
func NewParallelGetTask(url string, filePath string, concurrent int, tempFolder string) *ParallelGetTask {
	return &ParallelGetTask{
		// 参数赋值
		Url:        url,
		FilePath:   filePath,
		Concurrent: concurrent,
		TempFolder: tempFolder,
		// 初始化状态量
		TotalSize:     0,
		ShardTaskList: make([]*ShardTask, 0),
	}
}

// 发送HEAD请求获取待下载文件的大小
func (task *ParallelGetTask) getLength() error {
	// 发送请求
	response, e := http.Head(task.Url)
	if e != nil {
		color.Red("发送HEAD请求出错！")
		return e
	}
	// 读取并设定长度
	task.TotalSize = response.ContentLength
	return nil
}

// 根据待下载文件的大小和设定的并发数，创建每个分片任务对象
func (task *ParallelGetTask) allocateTask() {
	// 如果并发数大于总大小，则进行调整
	if int64(task.Concurrent) > task.TotalSize {
		task.Concurrent = int(task.TotalSize)
	}
	// 开始计算每个分片的下载范围
	eachSize := task.TotalSize / int64(task.Concurrent)
	// 创建任务对象
	for i := 0; i < task.Concurrent; i++ {
		task.ShardTaskList = append(task.ShardTaskList, NewShardTask(task.Url, i+1, filepath.Join(task.TempFolder, strconv.Itoa(i+1)), int64(i)*eachSize, int64(i+1)*eachSize-1))
	}
	// 处理末尾部分
	if task.TotalSize%int64(task.Concurrent) != 0 {
		task.ShardTaskList[task.Concurrent-1].RangeEnd = task.TotalSize - 1
	}
}

// 根据任务列表进行多线程分片下载操作
func (task *ParallelGetTask) downloadShard() {
	// 创建线程组
	waitGroup := &sync.WaitGroup{}
	// 开始执行全部分片下载线程
	for _, task := range task.ShardTaskList {
		go task.DoShardGet(waitGroup)
		waitGroup.Add(1)
	}
	// 等待全部下载完成
	waitGroup.Wait()
}

// 下载完成后，合并分片文件
func (task *ParallelGetTask) mergeFile() error {
	// 创建目的文件
	targetFile, e := os.OpenFile(task.FilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0755)
	if e != nil {
		color.Red("创建目标文件出错！")
		return e
	}
	// 创建写入器
	writer := bufio.NewWriter(targetFile)
	// 准备读取每个分片文件
	for _, shard := range task.ShardTaskList {
		shardFile, e := os.OpenFile(shard.ShardFilePath, os.O_RDONLY, 0755)
		if e != nil {
			color.Red("读取分片文件出错！")
			return e
		}
		reader := bufio.NewReader(shardFile)
		readBuffer := make([]byte, 1024*1024)
		for {
			// 读取每个分片文件，一次读取1KB
			readSize, readError := reader.Read(readBuffer)
			// 处理结束或错误
			if readError != nil {
				if readError == io.EOF {
					break
				} else {
					color.Red("读取分片文件出错！")
					return readError
				}
			}
			// 写入到最终合并的文件
			_, writeError := writer.Write(readBuffer[0:readSize])
			if writeError != nil {
				color.Red("写入合并文件出错！")
				return writeError
			}
			_ = writer.Flush()
		}
		// 关闭分片文件资源
		_ = shardFile.Close()
	}
	// 关闭目的文件资源
	_ = targetFile.Close()
	return nil
}

// 删除分片临时文件
func (task *ParallelGetTask) cleanShard() error {
	for _, shard := range task.ShardTaskList {
		e := os.Remove(shard.ShardFilePath)
		if e != nil {
			color.Red("删除分片临时文件%s出错！", shard.ShardFilePath)
			return e
		}
	}
	return nil
}

// 在一个新线程中，实时输出每个分片的下载进度和总进度
func (task *ParallelGetTask) printTotalProcess() {
	go func() {
		// 上一次统计时的已下载大小，用于计算速度
		var lastDownloadSize int64 = 0
		for {
			// 如果全部任务完成则结束输出，并统计并发数
			allDone := true
			// 当前并发数
			currentTaskCount := 0
			for _, shardTask := range task.ShardTaskList {
				if !shardTask.TaskDone {
					allDone = false
					currentTaskCount += 1
				}
			}
			if allDone {
				break
			}
			// 统计所有分片已下载大小之和
			var totalDownloadSize int64 = 0
			for _, shardTask := range task.ShardTaskList {
				totalDownloadSize += shardTask.DownloadSize
			}
			// 计算速度
			currentDownload := totalDownloadSize - lastDownloadSize
			lastDownloadSize = totalDownloadSize
			speedString := util.ComputeSpeed(currentDownload, 300)
			// 输出到控制台
			fmt.Printf("\r当前并发数：%3d 速度：%s 总进度：%3.2f%%", currentTaskCount, speedString, float32(totalDownloadSize)/float32(task.TotalSize)*100)
			// 等待300ms
			time.Sleep(300 * time.Millisecond)
		}
	}()
}

// Run 开始执行整个分片多线程下载任务
func (task *ParallelGetTask) Run() error {
	// 获取文件大小
	e := task.getLength()
	if e != nil {
		color.Red("%s", e)
		return e
	}
	color.HiYellow("已获取到下载文件大小：%d字节", task.TotalSize)
	// 分配任务
	task.allocateTask()
	color.HiYellow("已完成分片任务分配，共计%d个任务", len(task.ShardTaskList))
	// 开启进度输出
	task.printTotalProcess()
	// 开始下载分片
	task.downloadShard()
	color.HiYellow("\n所有分片已下载完成！")
	// 开始合并文件
	e = task.mergeFile()
	if e != nil {
		color.Red("%s", e)
		return e
	}
	color.HiYellow("合并分片完成！")
	// 清理临时分片文件
	e = task.cleanShard()
	if e != nil {
		color.Red("%s", e)
		return e
	}
	color.HiYellow("清理分片临时文件完成！")
	color.Green("分片下载任务完成！")
	return nil
}